from mmengine.config import read_base
with read_base():
    from .mgam import *

from mgamdata.models.MedNeXt import MM_MedNext_Encoder, MM_MedNext_Decoder_3D
from mgamdata.models.MoCo import MoCoV3, MoCoV3Head_WithAcc
from mmpretrain.models.losses import CrossEntropyLoss

embed_dims = 128
block_counts = [2,2,2,2,2,2,2,2,2]
MedNeXt_checkpoint = False  # torch.checkpoint
MedNext_KernalSize = 5

model = dict(
    type = MoCoV3,
    backbone_checkpoint=True,
    encoder=dict(
        type=MM_MedNext_Encoder,
        in_channels=in_channels, # type: ignore
        embed_dims=embed_dims,
        kernel_size=MedNext_KernalSize,
        exp_r=2,
        block_counts = block_counts,
        dim="3d",
        use_checkpoint=MedNeXt_checkpoint,
        norm_type='group',
    ),
    neck=None,
    decoder=dict(
        type=MM_MedNext_Decoder_3D,
        embed_dims=embed_dims,
        kernel_size=MedNext_KernalSize,
        exp_r=2,
        block_counts = block_counts,
        num_classes=embed_dims,
        out_channels=embed_dims,
        use_checkpoint=MedNeXt_checkpoint,
        deep_supervision=True,
        norm_type='group',
        loss_decode=dict(type=CrossEntropyLoss), 
    ),
    head=dict(
        type=MoCoV3Head_WithAcc, 
        loss=dict(type=CrossEntropyLoss), 
        embed_dim=embed_dims, 
        proj_channel=embed_dims,
        dim='3d',
    )
)
